package config

import (
    "ebbflow/src/crypto"
    "encoding/base64"
    "encoding/json"
    "fmt"
    "math/big"
    "os"
)

// JSON models for Paillier configs

type PaillierPublicKeyJSON struct {
    N string `json:"n_b64"` // base64 of big-endian bytes
    G string `json:"g_b64"`
}

type PaillierPrivateKeyJSON struct {
    Lambda string `json:"lambda_b64"`
    Mu     string `json:"mu_b64"`
}

type NodePaillierPub struct {
    NodeID int                  `json:"node_id"`
    PK     PaillierPublicKeyJSON `json:"pk"`
}

type NodePaillierPriv struct {
    NodeID int                   `json:"node_id"`
    SK     PaillierPrivateKeyJSON `json:"sk"`
}

type PaillierPubConfig struct {
    Keys []NodePaillierPub `json:"keys"`
}

type PaillierPrivConfig struct {
    Keys []NodePaillierPriv `json:"keys"`
}

// Helpers

func b64ToBigInt(s string) (*big.Int, error) {
    b, err := base64.StdEncoding.DecodeString(s)
    if err != nil { return nil, err }
    return new(big.Int).SetBytes(b), nil
}

func bigIntToB64(i *big.Int) string {
    return base64.StdEncoding.EncodeToString(i.Bytes())
}

// SavePaillierPublicConfig saves public keys to a JSON file
func SavePaillierPublicConfig(path string, m map[int]*crypto.PaillierPublicKey) error {
    cfg := PaillierPubConfig{Keys: make([]NodePaillierPub, 0, len(m))}
    for id, pk := range m {
        cfg.Keys = append(cfg.Keys, NodePaillierPub{
            NodeID: id,
            PK: PaillierPublicKeyJSON{N: bigIntToB64(pk.N), G: bigIntToB64(pk.G)},
        })
    }
    data, err := json.MarshalIndent(cfg, "", "  ")
    if err != nil { return err }
    return os.WriteFile(path, data, 0644)
}

// LoadPaillierPublicConfig loads public keys from a JSON file
func LoadPaillierPublicConfig(path string) (map[int]*crypto.PaillierPublicKey, error) {
    data, err := os.ReadFile(path)
    if err != nil { return nil, err }
    var cfg PaillierPubConfig
    if err := json.Unmarshal(data, &cfg); err != nil { return nil, err }
    out := make(map[int]*crypto.PaillierPublicKey)
    for _, e := range cfg.Keys {
        n, err := b64ToBigInt(e.PK.N); if err != nil { return nil, fmt.Errorf("node %d n: %w", e.NodeID, err) }
        g, err := b64ToBigInt(e.PK.G); if err != nil { return nil, fmt.Errorf("node %d g: %w", e.NodeID, err) }
        out[e.NodeID] = &crypto.PaillierPublicKey{N: n, NSquare: new(big.Int).Mul(n, n), G: g}
    }
    return out, nil
}

// SavePaillierPrivateConfig saves private keys to JSON (for secure local storage)
func SavePaillierPrivateConfig(path string, m map[int]*crypto.PaillierPrivateKey) error {
    cfg := PaillierPrivConfig{Keys: make([]NodePaillierPriv, 0, len(m))}
    for id, sk := range m {
        cfg.Keys = append(cfg.Keys, NodePaillierPriv{
            NodeID: id,
            SK: PaillierPrivateKeyJSON{Lambda: bigIntToB64(sk.Lambda), Mu: bigIntToB64(sk.Mu)},
        })
    }
    data, err := json.MarshalIndent(cfg, "", "  ")
    if err != nil { return err }
    return os.WriteFile(path, data, 0600)
}

// LoadPaillierPrivateConfig loads private keys from JSON (caller should only register local pid)
func LoadPaillierPrivateConfig(path string, pub map[int]*crypto.PaillierPublicKey) (map[int]*crypto.PaillierPrivateKey, error) {
    data, err := os.ReadFile(path)
    if err != nil { return nil, err }
    var cfg PaillierPrivConfig
    if err := json.Unmarshal(data, &cfg); err != nil { return nil, err }
    out := make(map[int]*crypto.PaillierPrivateKey)
    for _, e := range cfg.Keys {
        p := pub[e.NodeID]
        if p == nil { return nil, fmt.Errorf("missing public key for node %d", e.NodeID) }
        lambda, err := b64ToBigInt(e.SK.Lambda); if err != nil { return nil, err }
        mu, err := b64ToBigInt(e.SK.Mu); if err != nil { return nil, err }
        out[e.NodeID] = &crypto.PaillierPrivateKey{Lambda: lambda, Mu: mu, Public: p}
    }
    return out, nil
}







